import os
import json
import ollama
from pymilvus import MilvusClient
# 其他项目文件引用
from src import config

client = MilvusClient(
    uri=config.milvus_url,
    db_name="its",
)


def embedding(query):
    llm_client = ollama.Client(host=f"http://{config.ollama_host}:{config.ollama_port}")
    v = llm_client.embeddings(
        model="bge-m3:latest",
        prompt=query
    )
    return v.embedding


def query_from_milvus_by_vector(collection_name, vector, limit, output_fields):
    res = client.search(
        collection_name=collection_name,
        data=[vector],
        limit=limit,
        output_fields=output_fields)
    return list(res)[0]


def query_question_pk(text):
    return query_from_milvus_by_vector(config.milvus_question_pk_collection, embedding(text), config.milvus_limit,["pk"])


def get_question_mapping():
    root_dir = os.path.dirname(os.path.dirname(__file__))
    question_mapping_path = os.path.join(root_dir, 'resources', 'question_mapping.json')

    with open(question_mapping_path, "r", encoding="utf-8") as f:
        question_mapping = json.load(f)

    return question_mapping
